# 导包 
import torch
import PIL.Image as Image 
import torch.nn as nn
import torch.nn.functional as F 
import torchvision as tv
import numpy as np
import pandas as pd
import sklearn 
from torch.utils.data import DataLoader,Dataset 


hyper_parameter = {
    "bs":128,
    "epochs":10,
    "lr":1e-3,
}

# 写dataset  加载数据集啥的 
class customDs(Dataset):
    
    def __init__(self):
        super().__init__()
        pass
    
    def __len__(self):
        return 0

    def __getitem__(self,idx):
        # 这里需要将数据转为tensor 
        pass

train_ds = customDs()
test_ds = customDs()

train_dl = DataLoader(train_ds,batch_size=hyper_parameter['bs'],shuffle=True,drop_last=False)
test_dl  = DataLoader(test_ds,batch_size=hyper_parameter['bs'])

# 写模型 

class customModel(nn.Module):

    def __init__(self,):
        super().__init__()
        pass

    def forward(self,x):
        pass
        return x

#保存模型加载模型

# 定义损失函数 
def loss_function(pred,target):
    # 需要考虑是否为不可导 
    return 0

# 定义评价函数 

def evaluation(pred,target):
    pass
    return 0

model = customModel()
loss_fn = loss_function
opt = torch.optim.Adam(model.parameters(),lr=hyper_parameter['lr'])


model = model.cuda() 

def train(epoch):

    model.train()
    step = 0
    for x,y in train_dl:
        x = x.cuda()
        y = y.cuda()
        pred = model(x)
        loss = loss_fn(pred,y)
        score = evaluation(pred,y) # 不是每一个模型都有 
        if step % 30 == 0:# 每30步记一次  自己设定 
            print("train:   epoch: {} step: {} loss: {} score: {}".format(epoch,step,loss.item(),score))
        step += 1
        opt.zero_grad()
        loss.backward()
        opt.step()


def test():
    model.eval()
    with torch.no_grad():
        step  = 0
        for x,y in test_dl:
            x = x.cuda()
            y = y.cuda()
            pred = model(x)
            loss = loss_fn(pred,y)
            score = evaluation(pred,y)
            if step % 30 == 0:# 每30步记一次  自己设定 
                print("test: step: {} loss: {} score: {}".format(step,loss.item(),score))
            step += 1
    total_score = 0
    return total_score 


def fit(epochs):
    
    for epoch in range(epochs):
        train(epoch)

        if epoch % 1 == 0: # n epoch validation 一次
            test()

fit(hyper_parameter['epochs'])
        